package io.kiki.sba.registry.server.metadata.slot.status;

import io.kiki.sba.registry.common.model.metaserver.nodes.DataStoreNode;
import io.kiki.sba.registry.common.model.slot.Slots;
import io.kiki.sba.registry.common.model.slot.SlotsStatusResponse;
import io.kiki.sba.registry.server.metadata.lease.dataStoreNode.DataStoreNodeManager;
import io.kiki.sba.registry.server.metadata.monitor.SlotsMonitor;
import io.kiki.sba.registry.server.metadata.slot.SlotManager;
import io.kiki.sba.registry.server.metadata.slot.arrange.ScheduledSlotArranger;
import io.kiki.sba.registry.server.metadata.slot.balance.BalancePolicy;
import io.kiki.sba.registry.server.metadata.slot.balance.BalancePolicyImpl;
import io.kiki.sba.registry.server.shared.slot.SlotsUtils;
import io.kiki.sba.registry.util.MathUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Map;


public class SlotsStatusService {

    @Autowired
    private SlotManager slotManager;

    @Autowired
    private SlotsMonitor slotsMonitor;

    @Autowired
    private DataStoreNodeManager dataStoreNodeManager;

    @Autowired
    private ScheduledSlotArranger slotArranger;

    public SlotsStatusResponse getSlotsStatus() {
        boolean isStableSlots = slotsMonitor.isStableSlots();
        Slots slots = slotManager.getSlots();
        Map<String, Integer> leaderCounter = SlotsUtils.getSlotsLeaderCount(slots);
        Map<String, Integer> followerCounter = SlotsUtils.getSlotsSlotCount(slots);
        //
        boolean isLeaderSlotBalanced = isSlotsLeaderBalanced(leaderCounter, dataStoreNodeManager.getDataStoreNodeInfo().getList());
        boolean isFollowerSlotBalanced = isSlotsFollowerBalanced(followerCounter, dataStoreNodeManager.getDataStoreNodeInfo().getList());
        //
        return new SlotsStatusResponse(slots.getEpoch(), isLeaderSlotBalanced, isFollowerSlotBalanced, isStableSlots, slotArranger.isSlotsProtectionMode(), leaderCounter, followerCounter);
    }

    public boolean isSlotsLeaderBalanced(Map<String, Integer> leaderCounter, List<DataStoreNode> dataStoreNodeList) {

        if (CollectionUtils.isEmpty(dataStoreNodeList)) {
            return false;
        }

        BalancePolicy balancePolicy = new BalancePolicyImpl();
        int expectedLeaderTotal = slotManager.getSlotNums();
        if (leaderCounter.values().stream().mapToInt(Integer::intValue).sum() < expectedLeaderTotal) {
            return false;
        }
        int leaderHighAverage = MathUtils.divideCeil(expectedLeaderTotal, dataStoreNodeList.size());
        int leaderHighWaterMark = balancePolicy.getHighWaterMarkSlotLeaderNums(leaderHighAverage);

        for (DataStoreNode dataStoreNode : dataStoreNodeList) {
            String dataIp = dataStoreNode.getIp();
            if (leaderCounter.get(dataIp) == null) {
                return false;
            }
            int leaderCount = leaderCounter.getOrDefault(dataIp, 0);
            if (leaderCount > leaderHighWaterMark) {
                return false;
            }
        }
        return true;
    }

    public boolean isSlotsFollowerBalanced(Map<String, Integer> followerCounter, List<DataStoreNode> dataStoreNodeList) {

        if (CollectionUtils.isEmpty(dataStoreNodeList)) {
            return false;
        }
        BalancePolicy balancePolicy = new BalancePolicyImpl();
        int expectedFollowerTotal = slotManager.getSlotNums() * (slotManager.getSlotReplicaNums() - 1);
        if (slotManager.getSlotReplicaNums() < dataStoreNodeList.size()) {
            if (followerCounter.values().stream().mapToInt(Integer::intValue).sum() < expectedFollowerTotal) {
                return false;
            }
        }
        int followerHighAverage = MathUtils.divideCeil(expectedFollowerTotal, dataStoreNodeList.size());
        int followerHighWaterMark = balancePolicy.getHighWaterMarkSlotFollowerNums(followerHighAverage);

        for (DataStoreNode dataStoreNode : dataStoreNodeList) {
            String dataIp = dataStoreNode.getIp();
            int followerCount = followerCounter.getOrDefault(dataIp, 0);
            if (followerCount > followerHighWaterMark) {
                return false;
            }
        }
        return true;
    }


}
